/**********************************************************
 *
 *  affine_flow_cimp.cpp
 *
 *  The C++ mex implementation of affine_flow
 *
 *  Created by Dahua Lin, on April 4, 2012
 *
 **********************************************************/

#include <mex.h>

#include <cmath>


// Auxiliary computation function

/**
 * Decompose as 
 *   
 *   a = s + d 
 *   b = s - d
 *
 * So:
 *
 *   s = (a + b) / 2
 *   d = (a - b) / 2
 */
template<typename T>
inline void mean_diff(T a, T b, T& s, T& d)
{
    s = (a + b) * T(0.5);
    d = (a - b) * T(0.5);
}


template<typename T>
inline T cosh_sqrt(T x)  // cosh(sqrt(x))
{
    if (x == 0)
    {
        return T(1);
    }
    else if (x > 0)
    {
        return std::cosh(std::sqrt(x)); 
    }
    else
    {
        return std::cos(std::sqrt(-x));
    }
}

template<typename T>
inline T sinh_ratio_sqrt(T x)  // sinh(sqrt(x)) / sqrt(x)
{
    if (x == 0)
    {
        return T(1);
    }
    else if (x > 0)
    {
        T sx = std::sqrt(x);
        return std::sinh(sx) / sx;
    }
    else
    {
        T sx = std::sqrt(-x);
        return std::sin(sx) / sx;
    }
}

template<typename T>
inline T exp_deriv2(T x)  // (exp(x) - (1 + x)) / x^2
{
    if (std::abs(x) < 1e-5)
    {
        return 0.5 + x * (T(1)/6) + (x*x) * (T(1)/24);
    }
    else
    {
        return (std::exp(x) - (1 + x)) / (x * x);
    }
}



template<typename T>
inline void matvec2x2(const T *A, const T* x, T *y)
{
    y[0] = A[0] * x[0] + A[2] * x[1];
    y[1] = A[1] * x[0] + A[3] * x[1];
}


template<typename T>
inline void matmul2x2(const T* A, const T* B, T *C)
{
    C[0] = A[0] * B[0] + A[2] * B[1];
    C[1] = A[1] * B[0] + A[3] * B[1];
    C[2] = A[0] * B[2] + A[2] * B[3];
    C[3] = A[1] * B[2] + A[3] * B[3];
}

template<typename T>
inline void matinv2x2(const T *A, T *R, T det)
{
    T invdet = T(1) / det;
    
    R[0] = invdet * A[3];
    R[1] = -invdet * A[1];
    R[2] = -invdet * A[2];
    R[3] = invdet * A[0];
}



/**
 * A class capturing various quantities for the Lie algebraic
 * generator
 *
 * X = [a+v h-r x
 *      h+r a-v y]
 *
 */ 
template<typename T>        
class LieGen
{
public:
    LieGen(const T *X) : Xmat(X)
    {
        // basic quantities
        
        T X1 = X[1];   // h + r
        T X2 = X[2];   // h - r
        
        mean_diff(X[0], X[3], a, v);
        mean_diff(X1, X2, h, r);
        x = X[4];
        y = X[5];
        
        // dependent quantities
        
        d2 = X1 * X2 + v * v;      
        
        T detX = X[0] * X[3] - X1 * X2;
        
        if (std::abs(detX) > 1e-10)
        {
            is_degen = false;
            
            T invY[4];
            matinv2x2(X, invY, detX);
            matvec2x2(invY, X+4, invY_u);
        }
        else
        {
            is_degen = true;
        }
    }
    
public:
    void linear_transform_part(T t, T *A) const
    {
        T e_a = std::exp(a * t);
                
        T t2 = t * t;
        
        T c = cosh_sqrt(d2 * t2);
        T s = sinh_ratio_sqrt(d2 * t2);
        
        T s_t = s * t;
        T vs_t = v * s_t;
        
        A[0] = e_a * (c + vs_t);
        A[1] = e_a * (h + r) * s_t;
        A[2] = e_a * (h - r) * s_t;
        A[3] = e_a * (c - vs_t);        
    }   
    
public:
    const T *Xmat;
    T a;
    T v; 
    T h;
    T r;
    T x;
    T y;
    
    bool is_degen;  // whether det(Y) = 0;
        
    T d2;           // h^2 - r^2 + v^2
    T invY_u[2];    // inv(Y) * u  (only used when is_degen = false);
            
}; // end class LieGen



template<typename T>
void compute_transform(const LieGen<T>& G, T t, T *Tmat)
{
    // compute the linear transform part
    
    G.linear_transform_part(t, Tmat);
    
    // compute the translation part
    
    if (G.x == 0 && G.y == 0)
    {
        Tmat[4] = 0;
        Tmat[5] = 0;
    }
    else
    {   
        T *tv = Tmat + 4;
        
        if (G.is_degen)
        {
            const T *X = G.Xmat;
            const T *u = G.Xmat + 4;
            
            T g = exp_deriv2(2 * G.a * t) * (t*t);
            
            matvec2x2(X, u, tv);
            (tv[0] *= g) += (u[0] * t);
            (tv[1] *= g) += (u[1] * t);            
        }
        else
        {
            matvec2x2(Tmat, G.invY_u, tv);
            tv[0] -= G.invY_u[0];
            tv[1] -= G.invY_u[1];
        }
       
    }
}



template<typename T>
inline mxArray* do_compute(
        const mxArray *mxX, const mxArray *mxTs, mxArray *mxR, int n)
{
    const T *X = (const T*)mxGetData(mxX);
    const T *ts = (const T*)mxGetData(mxTs);
    T *Ts = (T*)mxGetData(mxR);
    
    LieGen<T> G(X);
    
    for (int i = 0; i < n; ++i, Ts += 6)
    {
        compute_transform(G, ts[i], Ts);        
    }    
}


/**
 * Inputs
 *   [0] X:     the 2 x 3 Lie algebraic matrix
 *   [1] ts:    the vector of time steps
 *
 * Output
 *   [0] Ts:    The sequence of transform matrices
 *
 */
void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
    // take inputs
    
    const mxArray *mxX = prhs[0];
    const mxArray *mxTs = prhs[1];

    // prepare output
    
    int n = mxGetNumberOfElements(mxTs);
    mxClassID cid = mxGetClassID(mxX);
    
    mxArray *mxR = NULL;
    if (n == 1)
    {
        mxR = mxCreateNumericMatrix(2, 3, cid, mxREAL);
    }
    else
    {
        mwSize dims[3] = {2, 3, n};       
        mxR = mxCreateNumericArray(3, dims, cid, mxREAL);
    }
    
    
    // main   
    
    switch (mxGetClassID(mxX))
    {    
        case mxDOUBLE_CLASS:                        
            do_compute<double>(mxX, mxTs, mxR, n);
            break;
            
        case mxSINGLE_CLASS:
            do_compute<float>(mxX, mxTs, mxR, n);
            break;
    }
    
    plhs[0] = mxR;
}



 